feat: distributed EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE (rebased onto spiceai-52.5) by phillipleblanc · Pull Request #34 · spiceai/datafusion-ballista

phillipleblanc · 2026-04-23T07:41:22Z

Reimplements #31 on top of spiceai-52.5 (DataFusion 52). Supersedes #31.

What

End-to-end support for the three EXPLAIN forms over distributed Ballista:

EXPLAIN ... — adds a distributed_plan row alongside logical_plan and physical_plan.
EXPLAIN FORMAT TREE ... — physical plan rendered with box-drawing tree characters; logical_plan is omitted to match DataFusion's native behavior.
EXPLAIN ANALYZE ... — a single Plan with Metrics row containing the per-stage, per-operator metrics gathered from the scheduler after the job completes.

How

EXPLAIN / FORMAT TREE — ExplainFormat doesn't survive a datafusion-proto round-trip, so the client wraps LogicalPlan::Explain in a new BallistaExplainNode logical extension before sending to the scheduler. The scheduler unwraps it back into a native LogicalPlan::Explain (unwrap_ballista_explain in state/mod.rs) and then its existing physical-planning intercept replaces ExplainExec with a distributed-aware variant. extract_logical_and_physical_plans and construct_distributed_explain_exec are threaded with &ExplainFormat so Tree/Indent/PgJson/Graphviz are honored.

EXPLAIN ANALYZE —

Client (BallistaQueryPlanner): strips LogicalPlan::Analyze, runs the inner plan via DistributedQueryExec, and wraps it in a new DistributedExplainAnalyzeExec. DistributedQueryExec gains an Arc<Mutex<Option<String>>> job_id handle that is populated when the job is submitted; the parent reads it after the child stream drains and calls the new GetJobMetrics RPC.
Scheduler: new GetJobMetrics RPC (proto: GetJobMetricsParams/Result, JobStageMetrics, OperatorWithMetrics, OperatorMetric) walks the execution graph in the same pre-order DFS order as ballista_core::utils::collect_plan_metrics so the per-operator metrics line up with the rendered plan text. Falls back from the active-job cache to the saved completed-job graph (get_job_execution_graph) so the call still succeeds after succeed_job moves the graph out of active_job_cache.

Tests

ballista/client/tests/context_checks.rs:

should_execute_explain_query_correctly (was already there)
should_execute_explain_format_tree_query_correctly — checks two-row result (no logical_plan) and box-drawing characters in physical_plan.
should_execute_explain_analyze_query — sanitizes metric values to ... and checks both stages, both ShuffleWriterExec and AggregateExec, and that per-operator metrics=[...] appear.

All three cases run in both standalone and remote modes (the remote mode goes over the gRPC scheduler, exercising the codec round-trip and the GetJobMetrics RPC):

test supported::should_execute_explain_query_correctly::case_1_standalone ... ok
test supported::should_execute_explain_query_correctly::case_2_remote ... ok
test supported::should_execute_explain_format_tree_query_correctly::case_1_standalone ... ok
test supported::should_execute_explain_format_tree_query_correctly::case_2_remote ... ok
test supported::should_execute_explain_analyze_query::case_1_standalone ... ok
test supported::should_execute_explain_analyze_query::case_2_remote ... ok

Also manually validated against a local distributed cluster (scheduler + executor with mTLS) by issuing the queries through Spice's HTTP API:

EXPLAIN returns the expected three rows.
EXPLAIN FORMAT TREE returns two rows with tree-formatted physical plan.

Notes / follow-ups

This PR covers the ballista client API (BallistaQueryPlanner) path. Spice's /v1/queries async path submits the raw LogicalPlan::Analyze directly via scheduler.submit_job, bypassing BallistaQueryPlanner, and the scheduler currently only intercepts Explain, not Analyze. A follow-up is needed in either Spice or in this scheduler submit_job path to get EXPLAIN ANALYZE working over /v1/queries.

Reimplements PR #31 on top of spiceai-52.5 (DataFusion 52). EXPLAIN Round-trip the ExplainFormat through the client -> scheduler boundary by wrapping the LogicalPlan::Explain in a BallistaExplainNode logical extension before serialization. The scheduler unwraps it back to a native LogicalPlan::Explain so its existing physical-planning intercept can substitute a distributed-aware ExplainExec replacement. EXPLAIN FORMAT TREE Honored end-to-end by threading the ExplainFormat through extract_logical_and_physical_plans and construct_distributed_explain_exec in scheduler/state/distributed_explain.rs (Tree format omits the logical_plan row to match DataFusion's native behavior). EXPLAIN ANALYZE - Client (BallistaQueryPlanner): strips the LogicalPlan::Analyze and runs the inner plan via DistributedQueryExec, wrapped in a new DistributedExplainAnalyzeExec. After the child stream drains, the wrapper publishes the job_id (added Arc<Mutex<Option<String>>> handle on DistributedQueryExec) and calls the scheduler's GetJobMetrics RPC. - Scheduler: new GetJobMetrics RPC walks the execution graph in the same pre-order DFS order as ballista_core::utils::collect_plan_metrics so per-operator metrics line up with the rendered plan text. Falls back from the active-job cache to the saved completed-job graph so the call still succeeds after succeed_job moves the graph out of active_job_cache. Includes ballista/client tests covering all three forms in both standalone and remote modes.

Copilot

Pull request overview

Adds end-to-end distributed support for EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE in Ballista (DataFusion 52), ensuring explain-format survives the client→scheduler proto round-trip and enabling scheduler-sourced per-stage/operator metrics for analyze output.

Changes:

Introduces BallistaExplainNode logical extension + codec/proto support to preserve ExplainFormat across datafusion-proto, and unwraps it on the scheduler back into native LogicalPlan::Explain.
Threads ExplainFormat through distributed explain rendering, matching DataFusion behavior for FORMAT TREE (omit logical_plan row; tree-rendered physical_plan).
Adds GetJobMetrics scheduler gRPC RPC and a client-side DistributedExplainAnalyzeExec that fetches and renders per-stage/operator metrics as Plan with Metrics.

Reviewed changes

Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
ballista/scheduler/src/state/task_manager.rs	Makes execution-graph lookup available for gRPC metrics retrieval.
ballista/scheduler/src/state/mod.rs	Unwraps `BallistaExplainNode`, captures `ExplainFormat`, and threads it into distributed explain exec construction.
ballista/scheduler/src/state/distributed_explain.rs	Formats explain output according to `ExplainFormat` and matches DataFusion’s `FORMAT TREE` row behavior.
ballista/scheduler/src/scheduler_server/grpc.rs	Adds `GetJobMetrics` RPC and stage/operator metrics serialization using plan traversal.
ballista/executor/src/execution_loop.rs	Updates test SchedulerGrpc mock to include the new RPC.
ballista/core/src/serde/mod.rs	Extends Ballista logical extension codec to encode/decode `BallistaExplainNode` via Ballista protobuf.
ballista/core/src/serde/generated/ballista.rs	Regenerates protobuf bindings for `ExplainNode` and `GetJobMetrics` messages/RPC.
ballista/core/src/planner.rs	Wraps `EXPLAIN` for distribution and implements distributed `EXPLAIN ANALYZE` via `DistributedExplainAnalyzeExec`.
ballista/core/src/extension.rs	Defines `BallistaExplainNode` and stable string mapping for explain formats.
ballista/core/src/execution_plans/mod.rs	Registers the new `distributed_explain_analyze` execution plan module.
ballista/core/src/execution_plans/distributed_query.rs	Adds a shared `job_id` handle so parents can query scheduler metrics post-execution.
ballista/core/src/execution_plans/distributed_explain_analyze.rs	New exec that runs a distributed job, then fetches and renders metrics as a single-row `EXPLAIN ANALYZE` result.
ballista/core/proto/ballista.proto	Adds `LogicalPlanExplainNode` and `GetJobMetrics` RPC + message types.
ballista/client/tests/context_checks.rs	Adds integration coverage for `EXPLAIN FORMAT TREE` and `EXPLAIN ANALYZE` in standalone + remote modes.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 18 out of 19 changed files in this pull request and generated 1 comment.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 21 out of 23 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot AI review requested due to automatic review settings April 23, 2026 07:41

Copilot started reviewing on behalf of phillipleblanc April 23, 2026 07:41 View session

Copilot AI reviewed Apr 23, 2026

View reviewed changes

Comment thread ballista/core/src/execution_plans/distributed_explain_analyze.rs

Comment thread ballista/scheduler/src/scheduler_server/grpc.rs

Comment thread ballista/scheduler/src/state/task_manager.rs

sgrebnov added 2 commits April 23, 2026 12:16

Support for Tree formatting + tests

bc42c25

Fortmatting

4a0a6e6

Copilot AI review requested due to automatic review settings April 23, 2026 09:20

Copilot started reviewing on behalf of sgrebnov April 23, 2026 09:20 View session

More formatting

a140df6

sgrebnov force-pushed the phillip/explain-analyze-52.5 branch from f7eed42 to a140df6 Compare April 23, 2026 09:23

Copilot AI reviewed Apr 23, 2026

View reviewed changes

Comment thread ballista/core/src/execution_plans/distributed_explain_analyze.rs Outdated

sgrebnov added 3 commits April 23, 2026 12:33

Add insta snapshots for FORMAT TREE integration tests

538e10a

Lint

bf42918

Improve

d4458ec

Copilot AI review requested due to automatic review settings April 23, 2026 09:35

Copilot started reviewing on behalf of sgrebnov April 23, 2026 09:36 View session

Lint

a7d549e

Copilot AI reviewed Apr 23, 2026

View reviewed changes

Comment thread ballista/core/src/execution_plans/distributed_explain_analyze.rs

Comment thread ballista/scheduler/src/scheduler_server/grpc.rs

sgrebnov mentioned this pull request Apr 23, 2026

Update datafusion-ballista crate spiceai/spiceai#10488

Merged

sgrebnov approved these changes Apr 23, 2026

View reviewed changes

sgrebnov merged commit e6bcc66 into spiceai-52.5 Apr 23, 2026
30 checks passed

sgrebnov deleted the phillip/explain-analyze-52.5 branch April 23, 2026 10:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: distributed EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE (rebased onto spiceai-52.5)#34

feat: distributed EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE (rebased onto spiceai-52.5)#34
sgrebnov merged 8 commits into
spiceai-52.5from
phillip/explain-analyze-52.5

phillipleblanc commented Apr 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

phillipleblanc commented Apr 23, 2026

What

How

Tests

Notes / follow-ups

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants